"""
AFL++ AI 种子生成提示词配置
可以在这里修改生成种子的数量、格式和质量要求
"""

#TODO 我需要加上AFL++原始的空的种子，默认的，删除越界的种子类型，避免AI生成的时间过长和解析AI的响应过长时间

# 基础提示词模板
BASE_PROMPT_TEMPLATE = """
你是一个专业的模糊测试和安全研究员。请为以下程序生成高质量的AFL++初始种子文件。

{program_info}

程序分析：
- 检测到的危险函数: {dangerous_functions}
- 输入处理方法: {input_methods}
- 文件类型: {file_type}

请生成{seed_count}个精心设计的测试用例，涵盖以下类别：

{categories}

重要指导原则：
- 考虑程序的实际使用场景
- 针对检测到的危险函数设计特定测试用例
- 包含各种字符：空字节、换行符、转义序列、Unicode字符
- 对于命令行程序，考虑参数注入
- 对于文件处理程序，考虑畸形文件结构

输出要求：
- 每个测试用例单独一行
- 直接输出测试用例内容，不要添加编号、引号或解释
- 确保输出可以直接保存为AFL++种子文件
- 如果某些字符需要转义，请直接使用原始字符

现在请输出{seed_count}个测试用例：
"""

# 分类配置
CATEGORIES_CONFIG = {
    "c_program": {
        "categories": [
            "1. 正常输入（2个）：有效的标准输入，典型使用场景数据",
            "2. 边界测试（2个）：空输入、缓冲区边界值、数值边界",
            "3. 格式错误（2个）：格式字符串攻击、编码异常、错误数据结构",
            "4. 攻击载荷（2个）：缓冲区溢出、命令注入、路径遍历",
            "5. 随机变异（2个）：随机字节、混合攻击向量"
        ],
        "seed_count": 10  # 针对C程序，每种类型生成2个，共10个
    },
    "cpp_program": {
        "categories": [
            "1. 正常输入（2个）：有效的标准输入，典型使用场景数据",
            "2. 边界测试（2个）：空输入、缓冲区边界值、数值边界",
            "3. 面向对象错误（2个）：虚函数表破坏、对象状态异常",
            "4. 内存管理错误（2个）：new/delete不匹配、智能指针滥用",
            "5. 模板和异常（2个）：模板特化错误、异常安全漏洞"
        ],
        "seed_count": 10
    },
    "binary_program": {
        "categories": [
            "1. 正常输入（3个）：有效的标准输入",
            "2. 边界测试（3个）：各种边界条件",
            "3. 格式错误（2个）：畸形输入数据",
            "4. 攻击载荷（2个）：常见攻击模式"
        ],
        "seed_count": 10
    }
}


# 根据程序类型获取配置
def get_prompt_config(program_type):
    """根据程序类型获取提示词配置"""
    if program_type == "c":
        return CATEGORIES_CONFIG["c_program"]
    elif program_type == "cpp":
        return CATEGORIES_CONFIG["cpp_program"]
    else:
        return CATEGORIES_CONFIG["binary_program"]


def build_prompt(program_info, analysis, program_type="c"):
    """构建完整的提示词"""
    config = get_prompt_config(program_type)

    categories_text = "\n".join(config["categories"])
    seed_count = config["seed_count"]

    # 格式化危险函数和输入方法
    dangerous_functions = ', '.join(analysis['dangerous_functions']) if analysis[
        'dangerous_functions'] else '未检测到常见危险函数'
    input_methods = ', '.join(analysis['input_methods']) if analysis['input_methods'] else '未知输入方法'

    prompt = BASE_PROMPT_TEMPLATE.format(
        program_info=program_info,
        dangerous_functions=dangerous_functions,
        input_methods=input_methods,
        file_type=analysis['file_type'],
        categories=categories_text,
        seed_count=seed_count
    )

    return prompt, seed_count